"""
Docling文档处理器
"""
import logging
import json
import time
import psutil
import platform
from pathlib import Path
from typing import Dict, List, Optional, Any

from docling.document_converter import DocumentConverter, PdfFormatOption
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions

# 尝试导入torch来检测MPS设备
try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

logger = logging.getLogger(__name__)

class DoclingProcessor:
    """Docling文档处理器类"""

    def __init__(self, model_manager=None):
        """初始化处理器"""
        self.model_manager = model_manager
        self._converter = None
        self._log_system_info()
        self._initialize_converter()

    def _log_system_info(self):
        """记录系统信息和设备状态"""
        logger.info("=" * 60)
        logger.info("🖥️  系统信息和设备状态")
        logger.info("=" * 60)

        # 基本系统信息
        logger.info(f"操作系统: {platform.system()} {platform.release()}")
        logger.info(f"处理器: {platform.processor()}")
        logger.info(f"Python版本: {platform.python_version()}")

        # 内存信息
        memory = psutil.virtual_memory()
        logger.info(f"总内存: {memory.total / (1024**3):.1f} GB")
        logger.info(f"可用内存: {memory.available / (1024**3):.1f} GB")
        logger.info(f"内存使用率: {memory.percent:.1f}%")

        # PyTorch和设备信息
        if TORCH_AVAILABLE:
            logger.info(f"PyTorch版本: {torch.__version__}")

            # 检查MPS设备
            if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🚀 MPS (Metal Performance Shaders) 设备可用")
                logger.info("   - 这是Apple Silicon Mac的GPU加速")
                logger.info("   - 将用于加速AI模型推理")

                # 尝试创建MPS张量来测试
                try:
                    test_tensor = torch.randn(10, 10, device='mps')
                    logger.info("   - MPS设备测试: ✅ 成功")
                    del test_tensor
                except Exception as e:
                    logger.warning(f"   - MPS设备测试: ❌ 失败 - {e}")
            else:
                logger.info("⚠️  MPS设备不可用，将使用CPU")

            # 检查CUDA设备（虽然Mac不支持，但记录一下）
            if torch.cuda.is_available():
                logger.info(f"CUDA设备数量: {torch.cuda.device_count()}")
                for i in range(torch.cuda.device_count()):
                    logger.info(f"  - GPU {i}: {torch.cuda.get_device_name(i)}")
            else:
                logger.info("CUDA设备: 不可用")
        else:
            logger.warning("PyTorch未安装，无法检测GPU设备")

        logger.info("=" * 60)
    
    def _initialize_converter(self):
        """初始化Docling转换器"""
        try:
            logger.info("🔧 正在初始化Docling标准处理器...")
            logger.info("=" * 60)
            logger.info("📋 Docling模型配置详情:")
            logger.info("   🤖 处理器类型: 标准Docling处理器")
            logger.info("   📚 核心AI模型套件:")
            logger.info("      🧠 布局检测模型:")
            logger.info("         - 模型: DocLayNet (Microsoft)")
            logger.info("         - 功能: 文档布局分析和区域检测")
            logger.info("         - 识别: 标题、段落、表格、图片、公式等")
            logger.info("      📊 表格识别模型:")
            logger.info("         - 模型: TableFormer")
            logger.info("         - 功能: 表格结构识别和数据提取")
            logger.info("         - 识别: 表格边界、单元格、行列结构")
            logger.info("      🔤 OCR文字识别:")
            logger.info("         - 模型: EasyOCR + PaddleOCR")
            logger.info("         - 功能: 光学字符识别")
            logger.info("         - 支持: 多语言文字识别")
            logger.info("      🧮 公式识别模型:")
            logger.info("         - 模型: Nougat (Meta)")
            logger.info("         - 功能: 数学公式识别和LaTeX转换")
            logger.info("         - 识别: 复杂数学表达式")
            logger.info("      🖼️  图像分类模型:")
            logger.info("         - 模型: ResNet-based分类器")
            logger.info("         - 功能: 图像内容分类和描述")
            logger.info("         - 识别: 图表、照片、示意图等")

            logger.info("   🎯 处理功能配置:")
            logger.info("      ✅ 图片分类: 启用")
            logger.info("      ✅ 图片描述: 启用")
            logger.info("      ✅ 页面图片生成: 启用")
            logger.info("      ✅ 图片提取: 启用")
            logger.info("      ✅ 图片缩放: 2.0x (高清处理)")

            # 配置PDF处理选项 - 使用标准管道
            pdf_options = PdfPipelineOptions(
                # 启用图片分类
                do_picture_classification=True,
                # 启用图片描述
                do_picture_description=True,
                # 生成页面图片
                generate_page_images=True,
                # 生成图片
                generate_picture_images=True,
                # 图片缩放
                images_scale=2.0,
                # 模型路径配置 - 使用默认路径
                artifacts_path=None
            )
            
            logger.info("   📦 支持的文档格式:")
            logger.info("      📄 PDF: 完整布局分析 + OCR + 表格识别")
            logger.info("      📝 DOCX: Office文档结构解析")
            logger.info("      📊 PPTX: 演示文稿内容提取")
            logger.info("      📈 XLSX: 电子表格数据提取")
            logger.info("      🌐 HTML: 网页结构解析")
            logger.info("      🖼️  图片: OCR文字识别")
            logger.info("      📋 Markdown: 文本结构解析")

            # 创建文档转换器 - 使用标准PDF管道
            logger.info("🔄 创建文档转换器...")
            self._converter = DocumentConverter(
                allowed_formats=[
                    InputFormat.PDF,
                    InputFormat.DOCX,
                    InputFormat.PPTX,
                    InputFormat.XLSX,
                    InputFormat.HTML,
                    InputFormat.IMAGE,
                    InputFormat.MD,
                ],
                format_options={
                    InputFormat.PDF: PdfFormatOption(
                        pipeline_options=pdf_options
                    )
                }
            )

            logger.info("✅ Docling标准处理器初始化成功!")
            logger.info("🚀 AI模型已就绪，具备以下能力:")
            logger.info("   🧠 智能布局分析 (DocLayNet)")
            logger.info("   📊 精确表格识别 (TableFormer)")
            logger.info("   🔤 多语言OCR识别 (EasyOCR)")
            logger.info("   🧮 数学公式处理 (Nougat)")
            logger.info("   🖼️  图像内容分析 (ResNet)")
            logger.info("   ⚡ 性能特点: 快速、准确、低内存占用")
            logger.info("=" * 60)
            
        except Exception as e:
            logger.error(f"初始化Docling转换器失败: {str(e)}")
            raise
    
    def process_document(
        self, 
        file_path: Path, 
        output_formats: List[str] = None,
        max_pages: int = 50
    ) -> Dict[str, Any]:
        """
        处理文档
        
        Args:
            file_path: 文档文件路径
            output_formats: 输出格式列表 ['markdown', 'json', 'html']
            max_pages: 最大处理页数
            
        Returns:
            处理结果字典
        """
        if output_formats is None:
            output_formats = ['markdown']
        
        try:
            # 开始处理日志
            start_time = time.time()
            file_size = file_path.stat().st_size

            logger.info("🚀 开始文档处理")
            logger.info("=" * 50)
            logger.info(f"📄 文件名: {file_path.name}")
            logger.info(f"📁 文件路径: {file_path}")
            logger.info(f"📊 文件大小: {file_size / (1024*1024):.2f} MB")
            logger.info(f"📋 最大页数: {max_pages}")
            logger.info(f"🎯 输出格式: {', '.join(output_formats)}")

            # 检查文件是否存在
            if not file_path.exists():
                raise FileNotFoundError(f"文件不存在: {file_path}")

            # 记录处理前的系统状态
            memory_before = psutil.virtual_memory()
            logger.info(f"💾 处理前内存使用: {memory_before.percent:.1f}% ({memory_before.used / (1024**3):.1f} GB)")

            # 检测文件类型并显示将使用的模型
            file_extension = file_path.suffix.lower()
            logger.info("🤖 AI模型处理流程:")

            if file_extension == '.pdf':
                logger.info("   📄 PDF文档处理 - 启用完整AI模型套件:")
                logger.info("      🧠 DocLayNet: 页面布局分析")
                logger.info("      📊 TableFormer: 表格结构识别")
                logger.info("      🔤 EasyOCR: 文字光学识别")
                logger.info("      🧮 Nougat: 数学公式识别")
                logger.info("      🖼️  ResNet: 图像分类分析")
            elif file_extension in ['.docx', '.pptx', '.xlsx']:
                logger.info(f"   📝 Office文档处理 ({file_extension.upper()}):")
                logger.info("      📋 结构解析器: Office格式解析")
                logger.info("      🔤 文本提取器: 内容提取")
                logger.info("      📊 表格处理器: 数据结构化")
            elif file_extension in ['.png', '.jpg', '.jpeg', '.tiff', '.bmp']:
                logger.info("   🖼️  图像文档处理:")
                logger.info("      🔤 EasyOCR: 图像文字识别")
                logger.info("      🧠 布局分析: 图像区域检测")
                logger.info("      📊 表格识别: 图像中的表格")
            elif file_extension in ['.html', '.htm']:
                logger.info("   🌐 HTML文档处理:")
                logger.info("      📋 DOM解析器: 网页结构解析")
                logger.info("      🔤 文本提取器: 内容提取")
            else:
                logger.info("   📋 通用文档处理:")
                logger.info("      📝 文本解析器: 基础文本处理")

            # 检查MPS设备状态
            if TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🔥 AI模型加速设备: MPS (Metal Performance Shaders)")
                logger.info("   - GPU加速: 所有AI模型将使用Apple Silicon GPU")
                logger.info("   - 预期效果: 处理速度提升2-5倍")
                logger.info("   - 内存优化: GPU内存自动管理")
            else:
                logger.info("🖥️  AI模型运行设备: CPU")
                logger.info("   - 处理模式: CPU多线程处理")
                logger.info("   - 预期速度: 标准处理速度")

            logger.info("⏳ 开始AI模型处理...")
            logger.info("   🔄 正在加载和运行AI模型...")
            conversion_start = time.time()

            # 转换文档
            result = self._converter.convert(
                source=str(file_path),
                max_num_pages=max_pages
            )

            conversion_time = time.time() - conversion_start
            logger.info(f"✅ AI模型处理完成，耗时: {conversion_time:.2f} 秒")

            if not result.document:
                raise ValueError("文档转换失败，未生成文档对象")

            # 记录转换结果和模型处理效果
            pages_count = len(result.document.pages) if hasattr(result.document, 'pages') else 0
            logger.info(f"📖 AI模型处理结果: {pages_count} 页")

            # 分析处理的内容类型
            if hasattr(result.document, 'pages') and result.document.pages:
                logger.info("🔍 AI模型识别的内容类型:")
                text_blocks = 0
                table_blocks = 0
                image_blocks = 0

                for page in result.document.pages:
                    if hasattr(page, 'elements'):
                        for element in page.elements:
                            element_type = str(type(element).__name__).lower()
                            if 'text' in element_type or 'paragraph' in element_type:
                                text_blocks += 1
                            elif 'table' in element_type:
                                table_blocks += 1
                            elif 'image' in element_type or 'picture' in element_type:
                                image_blocks += 1

                logger.info(f"   📝 文本块: {text_blocks} (DocLayNet + OCR识别)")
                logger.info(f"   📊 表格: {table_blocks} (TableFormer识别)")
                logger.info(f"   🖼️  图像: {image_blocks} (ResNet分类)")

            # 记录处理后的内存状态
            memory_after = psutil.virtual_memory()
            memory_used = memory_after.used - memory_before.used
            logger.info(f"💾 AI模型内存使用: {memory_after.percent:.1f}% (增加: {memory_used / (1024**2):.1f} MB)")

            # 如果使用了MPS，记录相关信息
            if TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🔥 MPS GPU加速统计:")
                logger.info(f"   - AI模型处理速度: {pages_count / conversion_time:.2f} 页/秒")
                logger.info("   - GPU加速效果: 所有AI模型均已加速")
                logger.info("   - 模型推理: DocLayNet, TableFormer, OCR等均使用GPU")
            else:
                logger.info("🖥️  CPU处理统计:")
                logger.info(f"   - AI模型处理速度: {pages_count / conversion_time:.2f} 页/秒")
                logger.info("   - 处理模式: CPU多线程并行")
            
            # 生成不同格式的输出
            logger.info("📝 开始生成输出格式...")
            outputs = {}

            for format_name in output_formats:
                try:
                    format_start = time.time()
                    logger.info(f"   🔄 生成 {format_name.upper()} 格式...")

                    if format_name.lower() == 'markdown':
                        content = result.document.export_to_markdown()
                        outputs['markdown'] = content
                        logger.info(f"   ✅ Markdown生成完成 ({len(content)} 字符)")

                    elif format_name.lower() == 'json':
                        dict_data = result.document.export_to_dict()
                        content = json.dumps(dict_data, ensure_ascii=False, indent=2)
                        outputs['json'] = content
                        logger.info(f"   ✅ JSON生成完成 ({len(content)} 字符)")

                    elif format_name.lower() == 'html':
                        content = result.document.export_to_html()
                        outputs['html'] = content
                        logger.info(f"   ✅ HTML生成完成 ({len(content)} 字符)")

                    elif format_name.lower() == 'text':
                        content = result.document.export_to_markdown(strict_text=True)
                        outputs['text'] = content
                        logger.info(f"   ✅ Text生成完成 ({len(content)} 字符)")

                    else:
                        logger.warning(f"   ⚠️  不支持的输出格式: {format_name}")
                        continue

                    format_time = time.time() - format_start
                    logger.info(f"   ⏱️  {format_name.upper()}格式生成耗时: {format_time:.2f} 秒")

                except Exception as e:
                    logger.error(f"   ❌ 生成{format_name}格式时出错: {str(e)}")
                    outputs[format_name] = f"生成{format_name}格式时出错: {str(e)}"
            
            # 计算总处理时间
            total_time = time.time() - start_time

            # 添加详细的元数据
            metadata = {
                'filename': file_path.name,
                'file_size': file_path.stat().st_size,
                'file_size_mb': file_size / (1024*1024),
                'pages_processed': pages_count,
                'processing_status': 'success',
                'processing_time': total_time,
                'conversion_time': conversion_time,
                'average_speed': pages_count / total_time if total_time > 0 else 0,
                'output_formats': list(outputs.keys()),
                'device_used': 'MPS (GPU)' if (TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else 'CPU',
                'memory_usage_mb': memory_used / (1024**2) if 'memory_used' in locals() else 0
            }

            # 处理完成总结
            logger.info("🎉 文档处理完成!")
            logger.info("=" * 50)
            logger.info(f"📄 文件: {file_path.name}")
            logger.info(f"📊 总耗时: {total_time:.2f} 秒")
            logger.info(f"📖 处理页数: {pages_count}")
            logger.info(f"⚡ 平均速度: {pages_count / total_time:.2f} 页/秒")
            logger.info(f"💾 内存使用: {memory_used / (1024**2):.1f} MB" if 'memory_used' in locals() else "💾 内存使用: N/A")
            logger.info(f"🎯 输出格式: {', '.join(outputs.keys())}")
            logger.info(f"🖥️  处理设备: {metadata['device_used']}")

            if TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🔥 MPS加速效果: 处理速度显著提升")

            logger.info("=" * 50)

            return {
                'outputs': outputs,
                'metadata': metadata
            }
            
        except Exception as e:
            error_time = time.time() - start_time if 'start_time' in locals() else 0

            logger.error("❌ 文档处理失败!")
            logger.error("=" * 50)
            logger.error(f"📄 文件: {file_path.name if file_path else 'unknown'}")
            logger.error(f"⏱️  失败时间: {error_time:.2f} 秒")
            logger.error(f"🚨 错误信息: {str(e)}")
            logger.error(f"🔍 错误类型: {type(e).__name__}")

            # 记录系统状态
            if 'memory_before' in locals():
                memory_current = psutil.virtual_memory()
                logger.error(f"💾 当前内存使用: {memory_current.percent:.1f}%")

            # 如果是MPS相关错误，提供特殊提示
            if TORCH_AVAILABLE and 'mps' in str(e).lower():
                logger.error("🔥 MPS设备相关错误:")
                logger.error("   - 可能是GPU内存不足")
                logger.error("   - 建议减少处理页数或使用CPU模式")

            logger.error("=" * 50)

            return {
                'outputs': {},
                'metadata': {
                    'filename': file_path.name if file_path else 'unknown',
                    'processing_status': 'error',
                    'error_message': str(e),
                    'error_type': type(e).__name__,
                    'processing_time': error_time,
                    'device_used': 'MPS (GPU)' if (TORCH_AVAILABLE and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()) else 'CPU'
                }
            }
    
    def get_supported_formats(self) -> List[str]:
        """获取支持的文件格式"""
        return [
            'pdf', 'docx', 'pptx', 'xlsx', 'html', 'htm', 'md',
            'png', 'jpg', 'jpeg', 'tiff', 'bmp', 'gif', 'webp'
        ]
    
    def validate_file(self, file_path: Path) -> Dict[str, Any]:
        """
        验证文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            验证结果
        """
        try:
            if not file_path.exists():
                return {'valid': False, 'error': '文件不存在'}
            
            # 检查文件大小 (100MB限制)
            file_size = file_path.stat().st_size
            if file_size > 100 * 1024 * 1024:
                return {'valid': False, 'error': '文件太大，请选择小于100MB的文件'}
            
            # 检查文件扩展名
            file_ext = file_path.suffix.lower().lstrip('.')
            if file_ext not in self.get_supported_formats():
                return {'valid': False, 'error': f'不支持的文件格式: {file_ext}'}
            
            return {
                'valid': True, 
                'file_size': file_size,
                'file_format': file_ext
            }
            
        except Exception as e:
            return {'valid': False, 'error': f'文件验证失败: {str(e)}'}
    
    def get_processing_info(self) -> Dict[str, Any]:
        """获取处理器信息"""
        return {
            'supported_formats': self.get_supported_formats(),
            'max_file_size': '100MB',
            'available_outputs': ['markdown', 'json', 'html', 'text'],
            'converter_initialized': self._converter is not None
        }
